from ultralytics import YOLO
from typing import Union, List, Optional, Dict, Any
import cv2
import os

def yolo_detection(
    image_path: Union[str, List[str]],
    model_path: str = "./models/yolov10-elevator-best.pt",
    save_image: bool = True,
    output_dir: str = "output_files",
    conf_threshold: float = 0.25,
    iou_threshold: float = 0.7,
    search_class: List[str] = None,
    search_class_image_save: bool = False,
    crop_output_dir: str = "output_files/cropped_images"
) -> tuple:
    """
    YOLO目标检测函数，返回预测结果和可视化图像
    
    Args:
        image_path (Union[str, List[str]]): 单张图片路径或图片路径列表
        model_path (str): 模型文件路径，默认为"yolov10-elevator-best.pt"
        save_image (bool): 是否保存带预测框的图片，默认为True
        output_dir (str): 检测结果输出目录，默认为"output"
        conf_threshold (float): 置信度阈值，默认为0.25
        iou_threshold (float): IOU阈值，默认为0.7
        search_class (List[str]): 指定要搜索的类别列表（为空时返回所有类别）
        search_class_image_save (bool): 是否保存指定类别坐标裁剪出来的图片，默认为False
        crop_output_dir (str): 裁剪图片保存目录，默认为"cropped_images"
    
    Returns:
        tuple: (results, annotated_image, filtered_results) 
               - results: 原始预测结果(所有结果)
               - annotated_image: 带标注的图像
               - filtered_results: 过滤后的结果字典（指定类别的搜索结果），格式为 {cls_name: [detection_dict1, detection_dict2, ...]}
                  每个detection_dict包含:
                  - 'confidence': 置信度分数
                  - 'bbox': 边界框坐标 [x1, y1, x2, y2]
                  - 'class_id': 类别ID
                  - 'crop_path': 裁剪图片保存路径（如果search_class_image_save为True）
                  
    坐标含义:
        bbox坐标格式为 [x1, y1, x2, y2]，其中:
        - x1, y1: 边界框左上角坐标
        - x2, y2: 边界框右下角坐标
        坐标值为像素值，基于图像坐标系（左上角为原点(0,0)）
    """
    # 加载模型
    model = YOLO(model_path)
    
    # 读取原始图像用于裁剪
    original_image = None
    if search_class_image_save and isinstance(image_path, str):
        original_image = cv2.imread(image_path)
    
    # 进行预测
    results = model(
        image_path,
        conf=conf_threshold,
        iou=iou_threshold,
        save=save_image,
        project=output_dir if save_image else None
    )
    
    # 获取带标注的图像
    annotated_image = None
    if results:
        annotated_image = results[0].plot()  # 如果是单张图片
        
        # 如果需要保存到指定路径
        if save_image:
            os.makedirs(output_dir, exist_ok=True)
            output_path = os.path.join(output_dir, "detection_result.jpg")
            cv2.imwrite(output_path, annotated_image)
            print(f"检测结果已保存至: {output_path}")
    
    # 处理过滤结果
    filtered_results = {}
    if results and original_image is not None:
        # 创建裁剪图片保存目录
        if search_class_image_save:
            os.makedirs(crop_output_dir, exist_ok=True)
        
        for result_idx, result in enumerate(results):
            if result.boxes is not None:
                for box_idx, box in enumerate(result.boxes):
                    cls_id = int(box.cls)
                    cls_name = result.names[cls_id]
                    confidence = float(box.conf)
                    bbox = box.xyxy[0].tolist()  # [x1, y1, x2, y2]
                    
                    # 检查是否在指定类别中或search_class为空
                    if not search_class or cls_name in search_class:
                        detection_info = {
                            'confidence': confidence,
                            'bbox': bbox,
                            'class_id': cls_id
                        }
                        
                        # 保存裁剪图片
                        if search_class_image_save and original_image is not None:
                            # 转换坐标为整数
                            x1, y1, x2, y2 = map(int, bbox)
                            
                            # 确保坐标在图像范围内
                            height, width = original_image.shape[:2]
                            x1 = max(0, min(x1, width))
                            y1 = max(0, min(y1, height))
                            x2 = max(0, min(x2, width))
                            y2 = max(0, min(y2, height))
                            
                            # 裁剪图像
                            if x2 > x1 and y2 > y1:  # 确保有有效的区域
                                cropped_image = original_image[y1:y2, x1:x2]
                                
                                # 生成保存路径
                                base_name = os.path.splitext(os.path.basename(image_path))[0]
                                crop_filename = f"{base_name}_{cls_name}_{result_idx}_{box_idx}.jpg"
                                crop_path = os.path.join(crop_output_dir, crop_filename)
                                
                                # 保存裁剪图片
                                cv2.imwrite(crop_path, cropped_image)
                                detection_info['crop_path'] = crop_path
                        
                        if cls_name not in filtered_results:
                            filtered_results[cls_name] = []
                        filtered_results[cls_name].append(detection_info)
    
    return results, annotated_image, filtered_results

def yolo_detection_demo():
    # # 示例1: 基本使用（所有类别）
    # results, annotated_img, filtered_results = yolo_detection(
    #     image_path="elevator_images/ele-image.jpg",
    #     save_image=True,
    #     output_dir="detection_results"
    # )
    
    # print("所有检测结果:")
    # for cls_name, detections in filtered_results.items():
    #     print(f"类别 {cls_name}: {len(detections)} 个检测")
    #     for i, det in enumerate(detections):
    #         print(f"  检测{i+1}: 置信度={det['confidence']:.3f}, "
    #               f"坐标={[round(x) for x in det['bbox']]}")
    
    # 示例2: 检测特定类别并保存裁剪图片
    image_path="input_files/bus.jpg"#"input_files/elevator_images/ele-image.jpg"
    results, annotated_img, filtered_results = yolo_detection(
        image_path=image_path,
        search_class=["person", "1"],
        search_class_image_save=True,
        save_image=True,
    )
    
    print("\n检测person和elevator类别并保存裁剪图片:")
    for cls_name, detections in filtered_results.items():
        print(f"类别 {cls_name}: {len(detections)} 个检测")
        for i, det in enumerate(detections):
            print(f"  检测{i+1}: 置信度={det['confidence']:.3f}, "
                  f"坐标={[round(x) for x in det['bbox']]}")
            if 'crop_path' in det:
                print(f"    裁剪图片保存至: {det['crop_path']}")
    
    # # 示例3: 不保存整体图片，只保存裁剪图片
    # results, annotated_img, filtered_results = yolo_detection(
    #     image_path="elevator_images/ele-image.jpg",
    #     search_class=["person"],
    #     search_class_image_save=True,
    #     save_image=False,
    #     crop_output_dir="person_crops"
    # )

# 使用示例
if __name__ == "__main__":
    yolo_detection_demo()